# -*- coding: utf-8 -*-
import os
import re
import fitz
import pandas as pd
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm

# 预编译正则表达式
PAT_CLEAN = re.compile(r"\n|—+\d*—+")
PAT_KEYWORDS = re.compile(r"(.{0,10})(数字技术是指)(.{0,10})")

def save_dataframe(df):
    """优化存储逻辑"""
    if not os.path.exists(r"./result"):
        os.makedirs(r"./result")
    timestamp = pd.Timestamp.now().strftime("%Y%m%d_%H%M%S")
    output_path = f"./result/pdf_search_{timestamp}.xlsx"
    df.to_excel(output_path, index=False, engine="xlsxwriter")
        
def extract_page_text(page):
    """单页文本提取"""
    try:
        return PAT_CLEAN.sub("", page.get_text())
    except:
        return ""

def process_pdf(file_path):
    """并行处理单个PDF"""
    try:
        with fitz.open(file_path) as doc:
            full_text = (extract_page_text(page) for page in doc)
            matches = PAT_KEYWORDS.findall("".join(full_text))
            
        if matches:
            context = ["".join(m) for m in matches]
            return {
                "文件路径": os.path.dirname(file_path),
                "文件名": os.path.basename(file_path),
                "匹配内容": "\n".join(context),
                "频次": len(matches)
            }
    except Exception as e:
        print(f"处理失败 {file_path}: {str(e)}")
    return None

def os_walk_pdfs(target_dir,max_workers=4):
    """多线程PDF扫描"""
    file_list = [
        os.path.join(root, f) 
        for root, _, files in os.walk(target_dir) 
        for f in files if f.lower().endswith(".pdf")
    ]

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        results = list(tqdm(executor.map(process_pdf, file_list), total=len(file_list), desc="处理进度"))
    
    valid_data = [r for r in results if r]
    if valid_data:
        save_dataframe(pd.DataFrame(valid_data))

if __name__ == "__main__":
    target_dir = r"/media/sq/F:/A_Andrewz/B_Study/CSSCI/03_archive"  # 修改为实际路径
    os_walk_pdfs(target_dir)